#include <mutex>
#include <atomic>
#include <chrono>
#include <thread>
#include <fstream>
#include "acl/acl.h"
#include "yolov5_ascend.h"
#include "camera_stream.h"
#include <time.h>

#define MS std::chrono::milliseconds

int readLabels(char* txtPath, std::vector<std::string>& labels)
{
    std::ifstream fin(txtPath);
    if (!fin.is_open())
    {
        std::cout << "read label files failed!\n";
        return -1;
    }
    std::string label;
    while(getline(fin, label))
    {
        // 移除字符串中的非ASCII字符
        label.erase(std::remove_if(label.begin(), label.end(), [](char c) {
            return !(c >= 32 && c <= 126);  // ASCII 可打印字符范围
        }), label.end());
        labels.push_back(label);	
    }
    return 0;
}

void stressTest(int argc, char** argv)
{
    Yolov5AscendDetConfigs configs;
    configs.omPath = argv[2];
    int inputSize = atoi(argv[3]);
    configs.inputSize = cv::Size(inputSize, inputSize);
    configs.scoreThr = 0.25;
    configs.nmsThr = 0.3;
    configs.means = {0.0f, 0.0f, 0.0f};
    configs.norms = {1.0f, 1.0f, 1.0f};
    std::vector<std::string> labels;
    readLabels(argv[4], labels);
    configs.classNum = labels.size();
    configs.deviceId = atoi(argv[6]);
    bool show = atoi(argv[7]);
    std::cout << "labels info:\n"; 
    for (int i = 0; i < labels.size(); i++)
    {
        std::cout << i << ": " << labels[i] << std::endl;
    }
    Yolov5Ascend *pyolo = new Yolov5Ascend(configs);
    CameraStream camera(argv[5]);
    camera.openStream();
    camera.streamLoop();

    std::vector<DetectionTarget> objects;
    cv::Mat frame;
    size_t frameNum = 0;
    std::time_t startStamp = std::time(0);
    while (true)
    {
        if (camera.getFrame(frame) != 0)
        {
            continue;
        }
        if (!frame.data)
        {
            std::cout << "frame data empty\n";
            continue;
        }
        auto start = std::chrono::steady_clock::now();
        objects.clear();
        pyolo->detect(frame, objects);
        auto end = std::chrono::steady_clock::now();
        std::chrono::duration<double> elapsed = end - start;
        float timeMs = elapsed.count() * 1000;
        if (show)
        {
            pyolo->drawObject(frame, objects, labels);
            cv::imshow("result", frame);
            cv::waitKey(1);
        }
        auto lastTime = std::time(0) - startStamp;
        std::cout << "time seconds: " << lastTime << "\tframe num: " << frameNum;
        std::cout << "\tobjects num: " << (int)objects.size() << "\t inference time: " << timeMs << " ms" << std::endl;
        frameNum++;
    }
    delete pyolo;
    pyolo = NULL;
}

void streamDemo(int argc, char** argv)
{
    Yolov5AscendDetConfigs configs;
    configs.omPath = argv[2];
    int inputSize = atoi(argv[3]);
    configs.inputSize = cv::Size(inputSize, inputSize);
    configs.scoreThr = 0.25;
    configs.nmsThr = 0.3;
    configs.means = {0.0f, 0.0f, 0.0f};
    configs.norms = {1.0f, 1.0f, 1.0f};
    std::vector<std::string> labels;
    readLabels(argv[4], labels);
    configs.classNum = labels.size();
    configs.deviceId = atoi(argv[6]);
    bool show = atoi(argv[7]);
    bool save = atoi(argv[8]);
    std::cout << "labels info:\n"; 
    for (int i = 0; i < labels.size(); i++)
    {
        std::cout << i << ": " << labels[i] << std::endl;
    }
    Yolov5Ascend *pyolo = new Yolov5Ascend(configs);
    cv::VideoCapture cap(argv[5]);
    if (!cap.isOpened())
    {
        std::cout << "open stream failed!\n";
        return;
    }
    int frameCount = (int)cap.get(cv::CAP_PROP_FRAME_COUNT);
    cv::VideoWriter writer;
    if (save)
    {
        std::time_t stamp = std::time(0);
        std::string saveName = "result_" + std::to_string(stamp) + ".mp4";
        double fps = cap.get(cv::CAP_PROP_FPS);
        cv::Size saveSize = cv::Size((int)cap.get(cv::CAP_PROP_FRAME_WIDTH),
                                    (int)cap.get(cv::CAP_PROP_FRAME_HEIGHT));
        writer.open(saveName, cv::VideoWriter::fourcc('m','p','4','v'), fps, saveSize);
        std::cout << "save video to " << saveName << std::endl;
    }

    std::vector<DetectionTarget> objects;
    cv::Mat frame;
    size_t frameNum = 0;
    std::time_t startStamp = std::time(0);
    while (true)
    {
        cap >> frame;
        if (!frame.data)
        {
            std::cout << "read stream finished.\n";
            break;
        }
        auto start = std::chrono::steady_clock::now();
        objects.clear();
        pyolo->detect(frame, objects);
        auto end = std::chrono::steady_clock::now();
        std::chrono::duration<double> elapsed = end - start;
        float timeMs = elapsed.count() * 1000;
        pyolo->drawObject(frame, objects, labels);
        if (save)
        {
            writer.write(frame);
        }
        if (show)
        {
            cv::imshow("result", frame);
            cv::waitKey(1);
        }
        auto lastTime = std::time(0) - startStamp;
        std::cout << "time seconds: " << lastTime << "total frames: " << frameCount << "\tcurrent frame: " << frameNum;
        std::cout << "\tobjects num: " << (int)objects.size() << "\t inference time: " << timeMs << " ms" << std::endl;
        frameNum++;
    }
    delete pyolo;
    pyolo = NULL;
}

int main(int argc, char **argv)
{
    if (argc != 9)
    {
        fprintf(stderr, "usage: %s [demo] [model] [input_size] [class_input] [video_input] [device_id] [show] [save] ..\n", argv[0]);
        return -1;
    }
    std::string demo = argv[1];
    if (demo == "stream")
    {
        streamDemo(argc, argv);
    }
    else if (demo == "stress")
    {
        stressTest(argc, argv);
    }
    else
    {
        std::cout << "unrecognized demo type: " << argv[1] << std::endl;
    }
    return 0;
}
